/**
 * 
 */
package edu.umd.clip.lm.model.training;

import java.util.*;

import edu.berkeley.nlp.util.Pair;
import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.factors.Dictionary;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.questions.*;
import edu.umd.clip.lm.util.*;
import edu.umd.clip.lm.util.tree.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ExchangeAlgo {
	public static Pair<HashSet<Integer>, HashSet<Integer>> getExchangeWordSplit(
			Set<Integer> set, 
			HashMap<Integer, Long2IntMap> eventsByWord, 
			Set<Integer> exchangeHints) 
	{
		final int MAX_ITERATIONS = 100 * set.size();
		
		if (set.size() <= 2) {
			Iterator<Integer> i = set.iterator();
			HashSet<Integer> set1 = new HashSet<Integer>(1);
			set1.add(i.next());
			HashSet<Integer> set2 = new HashSet<Integer>(1);
			if (i.hasNext()) {
				set2.add(i.next());
			}
			return new Pair<HashSet<Integer>, HashSet<Integer>>(set1, set2);
		}
		
		Queue<Integer> queue1 = new LinkedList<Integer>();
		Queue<Integer> queue2 = new LinkedList<Integer>();
		Long2IntMap set1stats = new Long2IntMap();
		Long2IntMap set2stats = new Long2IntMap();

		ArrayList<Integer> words1 = new ArrayList<Integer>(set.size() / 2);
		ArrayList<Integer> words2 = new ArrayList<Integer>(set.size() / 2);
		
		if (exchangeHints != null) {
			for(Integer word : set) {
				Long2IntMap map = eventsByWord.get(word);
				if (exchangeHints.contains(word)) {
					addWordToSet(set1stats, map);
					words1.add(word);
				} else {
					addWordToSet(set2stats, map);
					words2.add(word);
				}
			}
		} else {
			//System.err.println("Exchange: no hint provided");
		}
		
		if (words1.size() == 0 || words2.size() == 0) {
			words1.clear();
			words2.clear();
			
			for(Iterator<Integer> i = set.iterator(); i.hasNext();) {
				Integer word = i.next();
				words1.add(word);
				addWordToSet(set1stats, eventsByWord.get(word));
				
				if (!i.hasNext()) break;
				word = i.next();
				words2.add(word);
				addWordToSet(set2stats, eventsByWord.get(word));
			}
		} else {
			//System.err.printf("Exchange: taking hint %d/%d\n", words1.size(), words2.size());
		}
		
		final boolean RANDOMIZE_ORDER = false;
		
		if (RANDOMIZE_ORDER) {
			// randomize word order
			RandomizedQueue<Integer> q1 = new RandomizedQueue<Integer>(words1);
			q1.drainTo(queue1);
			words1 = null;
			
			RandomizedQueue<Integer> q2 = new RandomizedQueue<Integer>(words2);
			q2.drainTo(queue2);
			words2 = null;
		} else {
			queue1.addAll(words1);
			queue2.addAll(words2);
			words1 = null;
			words2 = null;
		}
		
		HashSet<Integer> set1 = new HashSet<Integer>(queue1);
		HashSet<Integer> set2 = new HashSet<Integer>(queue2);
		
		long set1TotalCount = set1stats.sumValues();
		long set2TotalCount = set2stats.sumValues();
		
		/*
		double likelihood1 = computeSetLikelihood(set1stats);
		double likelihood2 = computeSetLikelihood(set2stats);
		
		System.err.printf("Initial Likelihoods: %e (%d) %e (%d)\n", likelihood1, set1TotalCount,
					likelihood2, set2TotalCount);
		*/
		//double bestLikelihood = computeSetLikelihood(set1stats) + computeSetLikelihood(set2stats);
		
		int totalIterations = 0;
		boolean set1updated = false;
		boolean set2updated = false;
		
		//System.err.printf("Starting Exchange algorithm: %d words\n", set.size());
		 
		while(true) {
			++totalIterations;
			if (totalIterations >= MAX_ITERATIONS) {
				/*
				System.err.printf("STOPPING the Exchange Algorithm at iteration %d! total: %d words, set1: %d, set2: %d, queue1: %d, queue2: %d\n", 
						totalIterations, set.size(), set1.size(), set2.size(), queue1.size(), queue2.size());
				*/
				break;
			}
			boolean done = true;
			if (queue1.size() > 0) {
				Integer word = queue1.poll();
				if (set1.size() == 1) {
					// don't try to remove the last word
					continue;
				}
				Long2IntMap wordCounts = eventsByWord.get(word); 
				
				long wordTotalCount = wordCounts.sumValues();
				
				double llDiff = computeLikelihoodDifferenceIfMoved(set1stats, set1TotalCount, set2stats, set2TotalCount, wordCounts, wordTotalCount);
				
				
				//double likelihood = computeSetLikelihood(set1stats) + computeSetLikelihood(set2stats);

				//assert(ProbMath.approxEqual(llDiff, likelihood - bestLikelihood));
				
				if (llDiff > 0) {
					// accept the move
					//bestLikelihood = likelihood;
					set1.remove(word);
					set2.add(word);
					set1updated = true;
					set2updated = true;
					set1TotalCount -= wordTotalCount;
					set2TotalCount += wordTotalCount;

					addWordToSet(set2stats, wordCounts);
					substractWordFromSet(set1stats, wordCounts);
				} else {
					// reject the move
				}
				done = false;
			} else if (set1updated) {
				queue1.addAll(set1);
				set1updated = false;
				done = false;
			}
			
			if (queue2.size() > 0) {
				Integer word = queue2.poll();
				if (set2.size() == 1) {
					// don't try to remove the last word
					continue;
				}
				Long2IntMap wordCounts = eventsByWord.get(word); 
				
				long wordTotalCount = wordCounts.sumValues();
				
				double llDiff = computeLikelihoodDifferenceIfMoved(set2stats, set2TotalCount, set1stats, set1TotalCount, wordCounts, wordTotalCount);

				//double likelihood = computeSetLikelihood(set1stats) + computeSetLikelihood(set2stats);

				//assert(ProbMath.approxEqual(llDiff, likelihood - bestLikelihood));

				if (llDiff > 0) {
					// accept the move
					//bestLikelihood = likelihood;
					set2.remove(word);
					set1.add(word);
					set1updated = true;
					set2updated = true;
					set1TotalCount += wordTotalCount;
					set2TotalCount -= wordTotalCount;

					addWordToSet(set1stats, wordCounts);
					substractWordFromSet(set2stats, wordCounts);
				} else {
					// reject the move
				}
				done = false;
			} else if (set2updated) {
				queue2.addAll(set2);
				set2updated = false;
				done = false;
			}
			if (done) break;
		}
		//System.err.printf("Exchange algorithm: %d iterations. %d words\n", totalIterations, set.size());

		return new Pair<HashSet<Integer>, HashSet<Integer>>(set1, set2);
	}

	public static void addWordToSet(Long2IntMap to, Long2IntMap from) {
		for(Long2IntMap.Iterator i = from.iterator(); i.hasNext(); ) {
			Long2IntMap.Entry e = i.next();
			
			to.addAndGet(e.getKey(), e.getValue());
		}
	}
	
	public static void substractWordFromSet(Long2IntMap from, Long2IntMap what) {
		for(Long2IntMap.Iterator i = what.iterator(); i.hasNext(); ) {
			Long2IntMap.Entry e = i.next();
			
			int result = from.addAndGet(e.getKey(), - e.getValue());
			if (result == 0) {
				from.remove(e.getKey());
			}
		}
	}
	
	public static double computeSetLikelihood(Long2IntMap set) {
		double likelihood = 0;
		long totalCount = 0;
			
		for(Long2IntMap.Iterator i = set.iterator(); i.hasNext(); ) {
			int val = i.next().getValue();
			if (val == 0) continue;
			totalCount += val;
			likelihood += val * ProbMath.log2(val);
		}
		likelihood -= totalCount * ProbMath.log2(totalCount);
		return likelihood;
	}
	
	private static double computeLikelihoodDifferenceIfMoved(
			Long2IntMap fromSet, long fromSetTotalCount, 
			Long2IntMap toSet, long toSetTotalCount, 
			Long2IntMap whatSet, long whatSetTotalCount) 
	{
		double llDiff = 0;
		for(Long2IntMap.Iterator i = whatSet.iterator(); i.hasNext(); ) {
			Long2IntMap.Entry e = i.next();
			long word = e.getKey();
			int count = e.getValue();
			
			{
				// the From Set
				int oldCount = fromSet.get(word);
				int newCount = oldCount - count;
				
				llDiff -= oldCount * ProbMath.log2(oldCount);
				if (newCount > 0) {
					llDiff += newCount * ProbMath.log2(newCount);
				}
			}
			
			{
				// the To Set
				int oldCount = toSet.get(word);
				int newCount = oldCount + count;
				
				if (oldCount > 0) {
					llDiff -= oldCount * ProbMath.log2(oldCount);
				}
				llDiff += newCount * ProbMath.log2(newCount);
				
			}
		}
		
		// total count
		llDiff += fromSetTotalCount * ProbMath.log2(fromSetTotalCount);
		llDiff -= (fromSetTotalCount - whatSetTotalCount) * ProbMath.log2(fromSetTotalCount - whatSetTotalCount);

		llDiff += toSetTotalCount * ProbMath.log2(toSetTotalCount);
		llDiff -= (toSetTotalCount + whatSetTotalCount) * ProbMath.log2(toSetTotalCount + whatSetTotalCount);
		
		return llDiff;
	}
	
	public static String wordListToString(byte wordIndex, Collection<Integer> words) {
		StringBuilder sb = new StringBuilder();
		edu.umd.clip.lm.factors.Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(wordIndex);
		Iterator<Integer> i = words.iterator();
		sb.append('[');
		if (i.hasNext()) {
			sb.append('\'');
			sb.append(dict.getWord(i.next()));
			while(i.hasNext()) {
				sb.append("','");
				sb.append(dict.getWord(i.next()));
			}
			sb.append('\'');
		}
		sb.append(']');
		return sb.toString();
	}
	
	public static Collection<Pair<Question,Boolean>> getWordQuestions(byte factorIdx, byte timeOffset, BinaryTree<HistoryTreePayload> node) {
		LinkedList<Pair<Question,Boolean>> questions = new LinkedList<Pair<Question,Boolean>>();
		
		BinaryTree<HistoryTreePayload> child = node;
		BinaryTree<HistoryTreePayload> tree = node.getParent();
		while(tree != null) {
			HistoryTreePayload payload = tree.getPayload();
			Question q = payload.question;
			if (q.getTimeOffset() == timeOffset && q.getIndex() == factorIdx && !q.isAboutHidden()) {
				boolean answer = tree.getRight() == child;
				questions.addFirst(new Pair<Question,Boolean>(payload.question, answer));
			}
			child = tree;
			tree = tree.getParent();
		}
		return questions;
	}
	

	public static Set<Integer> getPossibleFactors(byte factorIdx, Collection<Pair<Question,Boolean>> questions) {
		Experiment exp = Experiment.getInstance();
		/*
		if (factorIdx == exp.getTupleDescription().getMainFactorIndex()) {
			return getPossibleWords(exp.getWordTree().getTree().copy(), questions);
		}
		*/
		// normal overt factors, use vocabulary directly
		Dictionary dict = exp.getTupleDescription().getDictionary(factorIdx);
		HashSet<Integer> possibleFactors = new HashSet<Integer>(dict.size());
		for(DictionaryIterator it = dict.iterator(true); it.hasNext();) {
			possibleFactors.add(it.next());
		}
		
		for(Pair<Question, Boolean> pair : questions) {
			Question question = pair.getFirst();
			boolean answer = pair.getSecond();
			
			if (question.getQuestionType() == Question.IN_SET_QUESTION) {
				InSetQuestion q = (InSetQuestion) question;
				if (answer) {
					for(Iterator<Integer> it = possibleFactors.iterator(); it.hasNext();) {
						int value = it.next();
						if (!q.hasValue(value)) it.remove();
					}
				} else {
					for(Iterator<Integer> it = possibleFactors.iterator(); it.hasNext();) {
						int value = it.next();
						if (q.hasValue(value)) it.remove();
					}					
				}
			} else if (question.getQuestionType() == Question.IDENTITY_QUESTION) {
				IdentityQuestion q = (IdentityQuestion) question;
				if (answer) {
					if (possibleFactors.contains(q.getValue())) {
						possibleFactors.clear();
						possibleFactors.add(q.getValue());
					} else {
						possibleFactors.clear();
						return possibleFactors;
					}
				} else {
					possibleFactors.remove(q.getValue());
				}
			} else {
				System.err.printf("getPossibleFactors(%d) : unexpected question: %s\n", factorIdx, question.toString());
			}
		}
		return possibleFactors;
	}

}
